"""a more advanced example of basic_tree.py.  treenodes can now reference their "root" node, and
introduces a new selection method which selects an entire tree of nodes at once, taking 
advantage of a custom MapperExtension to assemble incoming nodes into their correct structure."""

import os
import zlib
import cPickle
import base64
import time
import operator
import datetime

from sqlite3 import dbapi2 as sqlite
from sqlalchemy import create_engine
from sqlalchemy import MetaData, Table, Column, ForeignKey
from sqlalchemy import Integer, String, DateTime
from sqlalchemy.orm import mapper, relation, create_session
from sqlalchemy.sql import select

from utilities import FormatFileSize
from constants import _knownPackages


PICKLE_PROTOCOL = cPickle.HIGHEST_PROTOCOL
COMPRESSION_LEVEL = zlib.Z_BEST_COMPRESSION

def serialize(obj):
    return base64.b64encode(zlib.compress(cPickle.dumps(obj, PICKLE_PROTOCOL), COMPRESSION_LEVEL))

def deserialize(obj):
    return cPickle.loads(zlib.decompress(base64.b64decode(obj)))


class SQLData(object):

    def __init__(self, value=None):

        self.id = None
        self.value = value
        

class SQLNode(object):
    
    def __init__(self, name, filename=None, id=None, scan_date=None):

        self.name = name
        self.filename = filename
        
        self.id = id
        self.data = None
        
        self.scan_date = scan_date
        self.size = 0
        self.python_version = ""
        self.num_children = 0
        
        self.icon = self.GetIcon()
        
        
    def SetData(self, data):

        self.data = data


    def GetData(self):

        return deserialize(self.data.value)
    

    def GetName(self):

        return self.name


    def GetFilename(self):

        return self.filename


    def GetScanDate(self):

        return self.scan_date


    def SetScanDate(self, scan_date):

        self.scan_date = scan_date


    def GetIcon(self):

        if self.icon is not None:
            return self.icon

        lower = self.name.lower()
        
        for pkg in _knownPackages:
            if pkg in lower:
                self.icon = pkg
                return self.icon

        return "python"


    def SetSizeOnDisk(self, size):

        self.size = int(size/1024.0)


    def GetSizeOnDisk(self):

        return self.size
    

class Database(object):

    def __init__(self, mainframe):

        self.MainFrame = mainframe
        self.dbTree = mainframe.dbTree
        self.inputPanel = mainframe.inputPanel


    def Load(self, dbName):

        self.dbName = dbName

        self.engine = create_engine('sqlite+pysqlite:///%s'%dbName, echo=False, module=sqlite)

        metadata = MetaData(self.engine)

        package = Table('packages', metadata,
                        Column('id', Integer, primary_key=True),
                        Column('parent_id', Integer, ForeignKey('data.data_id'), nullable=True),
                        Column('name', String(50)),
                        Column('filename', String(100)),
                        Column('scan_date', DateTime),
                        Column('icon', String(20)),
                        Column('size', Integer),
                        Column('num_children', Integer),
                        Column('python_version', String(100)))

        data = Table("data", metadata,
                     Column('data_id', Integer, primary_key=True),
                     Column('value', String(100), nullable=False)
                     )

        metadata.create_all(bind=self.engine)
        mapper(SQLNode, package, properties=dict(id=package.c.id,
                                                 name=package.c.name,
                                                 filename=package.c.filename,
                                                 scan_date=package.c.scan_date,
                                                 icon=package.c.icon,
                                                 size=package.c.size,
                                                 num_children=package.c.num_children,
                                                 python_version=package.c.python_version,
                                                 data=relation(mapper(SQLData, data, properties=dict(id=data.c.data_id)),
                                                               cascade="all, delete")))

        self.session = create_session(bind=self.engine, autoflush=True, autocommit=True)

        allDB = self.PopulateInputPanel()

        self.MainFrame.SendMessage("Message", "Database location is: %s"%os.path.normpath(self.dbName))

        if allDB:
            size = FormatFileSize(os.path.getsize(self.dbName))
            number = len(allDB)
            modified = int(os.path.getmtime(self.dbName))
            modified = datetime.datetime.fromtimestamp(modified)

            self.MainFrame.SendMessage("Message", "Database contains %d library trees occupying %s (last modified: %s)"%(number, size, modified))


    def PopulateInputPanel(self):

        s = select([SQLNode.name, SQLNode.filename, SQLNode.scan_date, SQLNode.icon])

        allDB = self.engine.execute(s).fetchall()
                
        allDB = sorted(allDB, key=operator.attrgetter("scan_date"))
        self.inputPanel.PopulateCombos(allDB)

        return allDB
    

    def DeletePackage(self, package_name, vacuum=True):

        node = self.session.query(SQLNode).filter(SQLNode.name==package_name).first()
        
        if node:
            self.session.delete(node)
        
            self.session.flush()        
            self.session.expunge_all()

            if vacuum:
                self.Vacuum()


    def AddPackage(self, package_class):

        self.DeletePackage(package_class.name)
        
        package_class.Save()

        dbsizepre = os.path.getsize(self.dbName)
        
        scan_date = datetime.datetime(*time.localtime()[0:6])
        node = SQLNode(package_class.name, package_class.filename, scan_date=scan_date)
        node.num_children = package_class.GetChildrenCount()
        node.python_version = package_class.GetPythonVersion()
        node.SetData(SQLData(serialize(package_class)))

        self.session.add(node)
        self.session.flush()
        
        dbsizepost = os.path.getsize(self.dbName)
        increase = dbsizepost - dbsizepre

        node.SetSizeOnDisk(increase)
        self.session.add(node)
        self.session.flush()

        self.PopulateInputPanel()        


    def LoadPackage(self, package_name):

        sql_package = self.session.query(SQLNode).filter(SQLNode.name==package_name).first()
        package_class = sql_package.GetData()
        return package_class
                

    def GetAllPackages(self):

        s = select([SQLNode.name, SQLNode.filename, SQLNode.python_version, SQLNode.num_children, SQLNode.scan_date, SQLNode.icon, SQLNode.size])
        allDB = self.engine.execute(s).fetchall()
        allDB = sorted(allDB, key=operator.attrgetter("name"))

        size = FormatFileSize(os.path.getsize(self.dbName))
        number = len(allDB)
        modified = int(os.path.getmtime(self.dbName))
        modified = datetime.datetime.fromtimestamp(modified)

        dbInfo = [os.path.normpath(self.dbName), size, number, modified]
        
        return allDB, dbInfo
    

    def Vacuum(self):

        self.engine.execute("VACUUM")


    def CloseSession(self):

        self.session.flush()
        self.session.expunge_all()
        self.session.close()
        del self.session

